"""
18_ontology_design.py
=====================
Healthcare Ontology Design with Neo4j

This module covers:
1. Medical ontology structure (ICD-11, RxNorm, SNOMED CT, LOINC)
2. Entity-Relationship modeling
3. Hierarchical taxonomies
4. Cross-ontology mappings
5. Semantic relationships

Author: Neo4j Tutorial
Date: January 2025
"""

from typing import Dict, List, Any, Optional
from dataclasses import dataclass, field
from enum import Enum
from neo4j import Driver
import sys
import os

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from src.utils.connection import get_driver


class OntologyType(Enum):
    """Standard medical ontologies."""
    ICD11 = "ICD-11"      # International Classification of Diseases
    RXNORM = "RxNorm"     # Medication vocabulary
    SNOMED_CT = "SNOMED-CT"  # Clinical terminology
    LOINC = "LOINC"       # Lab observations
    CPT = "CPT"           # Procedure codes
    MESH = "MeSH"         # Medical Subject Headings


@dataclass
class OntologyNode:
    """Base class for ontology nodes."""
    code: str
    name: str
    ontology: OntologyType
    description: Optional[str] = None
    synonyms: List[str] = field(default_factory=list)
    parent_code: Optional[str] = None


@dataclass
class ICD11Disease:
    """ICD-11 Disease classification."""
    code: str            # e.g., "E11.9"
    name: str            # e.g., "Type 2 diabetes mellitus without complications"
    chapter: str         # e.g., "Endocrine, nutritional and metabolic diseases"
    block: str           # e.g., "Diabetes mellitus"
    category: str        # e.g., "E11"
    parent_code: Optional[str] = None
    inclusions: List[str] = field(default_factory=list)
    exclusions: List[str] = field(default_factory=list)


@dataclass
class RxNormMedication:
    """RxNorm Medication classification."""
    rxcui: str           # RxNorm Concept Unique Identifier
    name: str            # Drug name
    tty: str             # Term type (SCD, SBD, IN, etc.)
    generic_name: str
    brand_names: List[str] = field(default_factory=list)
    drug_class: Optional[str] = None
    ingredients: List[str] = field(default_factory=list)
    dose_form: Optional[str] = None
    strength: Optional[str] = None


@dataclass
class SNOMEDConcept:
    """SNOMED CT Clinical concept."""
    concept_id: str      # SNOMED CT ID
    fsn: str             # Fully Specified Name
    preferred_term: str
    semantic_tag: str    # e.g., "disorder", "finding", "procedure"
    hierarchy: str       # e.g., "Clinical finding"
    is_a: List[str] = field(default_factory=list)  # Parent concepts


class OntologyManager:
    """
    Manages healthcare ontology in Neo4j.
    
    Ontology Structure:
    
    ┌─────────────────────────────────────────────────────────────────────┐
    │                    HEALTHCARE ONTOLOGY MODEL                        │
    ├─────────────────────────────────────────────────────────────────────┤
    │                                                                     │
    │  ICD-11 HIERARCHY                    RXNORM HIERARCHY              │
    │  ────────────────                    ────────────────               │
    │                                                                     │
    │  ┌─────────────┐                    ┌─────────────┐                │
    │  │  Chapter    │                    │  DrugClass  │                │
    │  │ (Endocrine) │                    │  (Biguanide)│                │
    │  └──────┬──────┘                    └──────┬──────┘                │
    │         │IS_PARENT_OF                      │CONTAINS               │
    │         ▼                                  ▼                        │
    │  ┌─────────────┐                    ┌─────────────┐                │
    │  │   Block     │                    │ Ingredient  │                │
    │  │ (Diabetes)  │                    │ (Metformin) │                │
    │  └──────┬──────┘                    └──────┬──────┘                │
    │         │IS_PARENT_OF                      │INGREDIENT_OF          │
    │         ▼                                  ▼                        │
    │  ┌─────────────┐                    ┌─────────────┐                │
    │  │  Category   │                    │    Drug     │                │
    │  │   (E11)     │◄───TREATS────────► │ (Metformin  │                │
    │  └──────┬──────┘                    │  500mg tab) │                │
    │         │IS_PARENT_OF               └─────────────┘                │
    │         ▼                                                          │
    │  ┌─────────────┐                    CROSS-ONTOLOGY MAPPING         │
    │  │   Disease   │                    ─────────────────────          │
    │  │  (E11.9)    │◄───MAPS_TO────────►SNOMED CT Concept             │
    │  └─────────────┘                                                   │
    │                                                                     │
    └─────────────────────────────────────────────────────────────────────┘
    """
    
    def __init__(self, driver: Driver):
        self.driver = driver
    
    # =========================================================
    # SCHEMA CREATION
    # =========================================================
    
    def create_ontology_schema(self) -> Dict[str, Any]:
        """Create indexes and constraints for ontology nodes."""
        queries = [
            # ICD-11 indexes
            "CREATE INDEX icd11_code IF NOT EXISTS FOR (d:ICD11Disease) ON (d.code)",
            "CREATE INDEX icd11_chapter IF NOT EXISTS FOR (c:ICD11Chapter) ON (c.code)",
            "CREATE INDEX icd11_block IF NOT EXISTS FOR (b:ICD11Block) ON (b.code)",
            
            # RxNorm indexes
            "CREATE INDEX rxnorm_rxcui IF NOT EXISTS FOR (m:RxNormDrug) ON (m.rxcui)",
            "CREATE INDEX rxnorm_name IF NOT EXISTS FOR (m:RxNormDrug) ON (m.name)",
            "CREATE INDEX ingredient_rxcui IF NOT EXISTS FOR (i:Ingredient) ON (i.rxcui)",
            
            # SNOMED indexes
            "CREATE INDEX snomed_id IF NOT EXISTS FOR (c:SNOMEDConcept) ON (c.conceptId)",
            
            # Full-text search
            """CREATE FULLTEXT INDEX disease_search IF NOT EXISTS 
               FOR (d:ICD11Disease) ON EACH [d.name, d.description]""",
            """CREATE FULLTEXT INDEX drug_search IF NOT EXISTS 
               FOR (m:RxNormDrug) ON EACH [m.name, m.genericName]"""
        ]
        
        results = []
        with self.driver.session() as session:
            for query in queries:
                try:
                    session.run(query)
                    results.append({"query": query[:50], "status": "success"})
                except Exception as e:
                    results.append({"query": query[:50], "status": "error", "error": str(e)})
        
        return {"indexes_created": len([r for r in results if r['status'] == 'success'])}
    
    # =========================================================
    # ICD-11 HIERARCHY
    # =========================================================
    
    def create_icd11_chapter(self, code: str, name: str, description: str = None) -> Dict:
        """Create an ICD-11 chapter (top level)."""
        query = """
        MERGE (c:ICD11Chapter:OntologyNode {code: $code})
        SET c.name = $name,
            c.description = $description,
            c.ontology = 'ICD-11',
            c.level = 'chapter',
            c.updatedAt = datetime()
        RETURN c
        """
        with self.driver.session() as session:
            result = session.run(query, code=code, name=name, description=description)
            record = result.single()
            return dict(record["c"]) if record else None
    
    def create_icd11_block(self, code: str, name: str, chapter_code: str) -> Dict:
        """Create an ICD-11 block and link to chapter."""
        query = """
        MATCH (chapter:ICD11Chapter {code: $chapter_code})
        MERGE (b:ICD11Block:OntologyNode {code: $code})
        SET b.name = $name,
            b.ontology = 'ICD-11',
            b.level = 'block',
            b.updatedAt = datetime()
        MERGE (chapter)-[:HAS_BLOCK]->(b)
        MERGE (b)-[:PART_OF]->(chapter)
        RETURN b
        """
        with self.driver.session() as session:
            result = session.run(query, code=code, name=name, chapter_code=chapter_code)
            record = result.single()
            return dict(record["b"]) if record else None
    
    def create_icd11_disease(self, disease: ICD11Disease) -> Dict:
        """Create an ICD-11 disease entity."""
        query = """
        MERGE (d:ICD11Disease:Disease:OntologyNode {code: $code})
        SET d.name = $name,
            d.chapter = $chapter,
            d.block = $block,
            d.category = $category,
            d.inclusions = $inclusions,
            d.exclusions = $exclusions,
            d.ontology = 'ICD-11',
            d.level = 'disease',
            d.updatedAt = datetime()
        
        WITH d
        OPTIONAL MATCH (parent:ICD11Block {code: $category})
        FOREACH (_ IN CASE WHEN parent IS NOT NULL THEN [1] ELSE [] END |
            MERGE (parent)-[:HAS_DISEASE]->(d)
            MERGE (d)-[:PART_OF]->(parent)
        )
        
        RETURN d
        """
        with self.driver.session() as session:
            result = session.run(
                query,
                code=disease.code,
                name=disease.name,
                chapter=disease.chapter,
                block=disease.block,
                category=disease.category,
                inclusions=disease.inclusions,
                exclusions=disease.exclusions
            )
            record = result.single()
            return dict(record["d"]) if record else None
    
    # =========================================================
    # RXNORM HIERARCHY
    # =========================================================
    
    def create_drug_class(self, code: str, name: str, description: str = None) -> Dict:
        """Create a drug class."""
        query = """
        MERGE (dc:DrugClass:OntologyNode {code: $code})
        SET dc.name = $name,
            dc.description = $description,
            dc.ontology = 'RxNorm',
            dc.updatedAt = datetime()
        RETURN dc
        """
        with self.driver.session() as session:
            result = session.run(query, code=code, name=name, description=description)
            record = result.single()
            return dict(record["dc"]) if record else None
    
    def create_ingredient(
        self,
        rxcui: str,
        name: str,
        drug_class_code: str = None
    ) -> Dict:
        """Create an ingredient (active substance)."""
        query = """
        MERGE (i:Ingredient:OntologyNode {rxcui: $rxcui})
        SET i.name = $name,
            i.ontology = 'RxNorm',
            i.tty = 'IN',
            i.updatedAt = datetime()
        
        WITH i
        OPTIONAL MATCH (dc:DrugClass {code: $dc_code})
        FOREACH (_ IN CASE WHEN dc IS NOT NULL THEN [1] ELSE [] END |
            MERGE (dc)-[:CONTAINS_INGREDIENT]->(i)
            MERGE (i)-[:BELONGS_TO_CLASS]->(dc)
        )
        
        RETURN i
        """
        with self.driver.session() as session:
            result = session.run(query, rxcui=rxcui, name=name, dc_code=drug_class_code)
            record = result.single()
            return dict(record["i"]) if record else None
    
    def create_rxnorm_drug(self, drug: RxNormMedication) -> Dict:
        """Create an RxNorm drug product."""
        query = """
        MERGE (m:RxNormDrug:Medication:OntologyNode {rxcui: $rxcui})
        SET m.name = $name,
            m.tty = $tty,
            m.genericName = $generic_name,
            m.brandNames = $brand_names,
            m.drugClass = $drug_class,
            m.doseForm = $dose_form,
            m.strength = $strength,
            m.ontology = 'RxNorm',
            m.updatedAt = datetime()
        
        WITH m
        UNWIND $ingredients AS ing_rxcui
        OPTIONAL MATCH (i:Ingredient {rxcui: ing_rxcui})
        FOREACH (_ IN CASE WHEN i IS NOT NULL THEN [1] ELSE [] END |
            MERGE (i)-[:INGREDIENT_OF]->(m)
            MERGE (m)-[:HAS_INGREDIENT]->(i)
        )
        
        RETURN m
        """
        with self.driver.session() as session:
            result = session.run(
                query,
                rxcui=drug.rxcui,
                name=drug.name,
                tty=drug.tty,
                generic_name=drug.generic_name,
                brand_names=drug.brand_names,
                drug_class=drug.drug_class,
                dose_form=drug.dose_form,
                strength=drug.strength,
                ingredients=drug.ingredients
            )
            record = result.single()
            return dict(record["m"]) if record else None
    
    # =========================================================
    # CROSS-ONTOLOGY MAPPING
    # =========================================================
    
    def create_ontology_mapping(
        self,
        source_code: str,
        source_ontology: str,
        target_code: str,
        target_ontology: str,
        mapping_type: str = "EQUIVALENT"
    ) -> bool:
        """
        Create mapping between ontologies.
        
        Mapping types:
        - EQUIVALENT: Exact match
        - BROADER: Source is more specific
        - NARROWER: Source is more general
        - RELATED: Related but not equivalent
        """
        query = """
        MATCH (source:OntologyNode)
        WHERE (source.code = $source_code OR source.rxcui = $source_code OR source.conceptId = $source_code)
          AND source.ontology = $source_ont
        MATCH (target:OntologyNode)
        WHERE (target.code = $target_code OR target.rxcui = $target_code OR target.conceptId = $target_code)
          AND target.ontology = $target_ont
        MERGE (source)-[r:MAPS_TO {type: $mapping_type}]->(target)
        SET r.updatedAt = datetime()
        RETURN source, target
        """
        with self.driver.session() as session:
            result = session.run(
                query,
                source_code=source_code,
                source_ont=source_ontology,
                target_code=target_code,
                target_ont=target_ontology,
                mapping_type=mapping_type
            )
            return result.single() is not None
    
    def create_treatment_relationship(
        self,
        drug_rxcui: str,
        disease_code: str,
        efficacy: str = "high",
        first_line: bool = False
    ) -> bool:
        """Create TREATS relationship between drug and disease."""
        query = """
        MATCH (m:RxNormDrug {rxcui: $rxcui})
        MATCH (d:ICD11Disease {code: $disease_code})
        MERGE (m)-[r:TREATS]->(d)
        SET r.efficacy = $efficacy,
            r.firstLine = $first_line,
            r.updatedAt = datetime()
        RETURN m, d
        """
        with self.driver.session() as session:
            result = session.run(
                query,
                rxcui=drug_rxcui,
                disease_code=disease_code,
                efficacy=efficacy,
                first_line=first_line
            )
            return result.single() is not None
    
    # =========================================================
    # QUERYING ONTOLOGY
    # =========================================================
    
    def get_disease_hierarchy(self, disease_code: str) -> Dict[str, Any]:
        """Get full hierarchy for a disease."""
        query = """
        MATCH (d:ICD11Disease {code: $code})
        OPTIONAL MATCH path = (d)-[:PART_OF*]->(ancestor)
        WITH d, collect(DISTINCT ancestor) as ancestors
        OPTIONAL MATCH (d)-[:MAPS_TO]->(mapped)
        RETURN 
            d.code as code,
            d.name as name,
            [a IN ancestors | {code: a.code, name: a.name, level: a.level}] as hierarchy,
            [m IN collect(mapped) | {code: coalesce(m.code, m.conceptId), name: m.name, ontology: m.ontology}] as mappings
        """
        with self.driver.session() as session:
            result = session.run(query, code=disease_code)
            record = result.single()
            return record.data() if record else None
    
    def get_drug_details(self, rxcui: str) -> Dict[str, Any]:
        """Get drug with ingredients and treatments."""
        query = """
        MATCH (m:RxNormDrug {rxcui: $rxcui})
        OPTIONAL MATCH (m)-[:HAS_INGREDIENT]->(i:Ingredient)
        OPTIONAL MATCH (m)-[:TREATS]->(d:ICD11Disease)
        OPTIONAL MATCH (i)-[:BELONGS_TO_CLASS]->(dc:DrugClass)
        RETURN 
            m.rxcui as rxcui,
            m.name as name,
            m.genericName as genericName,
            m.strength as strength,
            m.doseForm as doseForm,
            collect(DISTINCT {rxcui: i.rxcui, name: i.name}) as ingredients,
            collect(DISTINCT {code: d.code, name: d.name}) as treats,
            collect(DISTINCT dc.name) as drugClasses
        """
        with self.driver.session() as session:
            result = session.run(query, rxcui=rxcui)
            record = result.single()
            return record.data() if record else None
    
    def search_diseases(self, search_term: str, limit: int = 10) -> List[Dict]:
        """Search diseases using full-text index."""
        query = """
        CALL db.index.fulltext.queryNodes('disease_search', $term)
        YIELD node, score
        RETURN node.code as code, node.name as name, node.chapter as chapter, score
        ORDER BY score DESC
        LIMIT $limit
        """
        with self.driver.session() as session:
            result = session.run(query, term=search_term, limit=limit)
            return [record.data() for record in result]


# ============================================================
# DEMONSTRATION
# ============================================================

if __name__ == "__main__":
    print("=" * 60)
    print("Healthcare Ontology Design - Demonstration")
    print("=" * 60)
    
    print("""
    Medical Ontology Standards:
    
    1. ICD-11 (International Classification of Diseases)
       - Purpose: Diagnosis coding
       - Structure: Chapter → Block → Category → Disease
       - Example: E11.9 = Type 2 diabetes without complications
    
    2. RxNorm (Medication Vocabulary)
       - Purpose: Drug identification
       - Structure: DrugClass → Ingredient → Drug Product
       - Example: RXCUI 6809 = Metformin
    
    3. SNOMED CT (Clinical Terminology)
       - Purpose: Clinical documentation
       - Structure: Hierarchical concepts with relationships
       - Example: 73211009 = Diabetes mellitus
    
    4. LOINC (Lab Observations)
       - Purpose: Lab test identification
       - Example: 4548-4 = Hemoglobin A1c
    """)
    
    try:
        driver = get_driver()
        ontology = OntologyManager(driver)
        
        # Create schema
        print("\n📌 Creating Ontology Schema")
        print("-" * 40)
        result = ontology.create_ontology_schema()
        print(f"   Indexes created: {result['indexes_created']}")
        
        # Create ICD-11 hierarchy
        print("\n📌 Creating ICD-11 Hierarchy")
        print("-" * 40)
        
        # Chapter
        ontology.create_icd11_chapter(
            "05",
            "Endocrine, nutritional or metabolic diseases",
            "Diseases of the endocrine system and metabolism"
        )
        
        # Block
        ontology.create_icd11_block(
            "5A10-5A25",
            "Diabetes mellitus",
            "05"
        )
        
        # Disease
        disease = ICD11Disease(
            code="5A11",
            name="Type 2 diabetes mellitus",
            chapter="Endocrine, nutritional or metabolic diseases",
            block="Diabetes mellitus",
            category="5A10-5A25",
            inclusions=["adult-onset diabetes", "non-insulin-dependent diabetes"],
            exclusions=["Type 1 diabetes mellitus"]
        )
        ontology.create_icd11_disease(disease)
        print("   Created: Chapter → Block → Disease hierarchy")
        
        # Create RxNorm hierarchy
        print("\n📌 Creating RxNorm Hierarchy")
        print("-" * 40)
        
        # Drug class
        ontology.create_drug_class(
            "BIGUANIDE",
            "Biguanide",
            "Oral antidiabetic agents that decrease hepatic glucose production"
        )
        
        # Ingredient
        ontology.create_ingredient("6809", "Metformin", "BIGUANIDE")
        
        # Drug product
        drug = RxNormMedication(
            rxcui="861004",
            name="Metformin 500 MG Oral Tablet",
            tty="SCD",
            generic_name="Metformin",
            brand_names=["Glucophage"],
            drug_class="Biguanide",
            ingredients=["6809"],
            dose_form="Oral Tablet",
            strength="500 MG"
        )
        ontology.create_rxnorm_drug(drug)
        print("   Created: DrugClass → Ingredient → Drug hierarchy")
        
        # Create treatment relationship
        print("\n📌 Creating Treatment Relationship")
        print("-" * 40)
        ontology.create_treatment_relationship(
            "861004", "5A11",
            efficacy="high",
            first_line=True
        )
        print("   Metformin 500mg TREATS Type 2 Diabetes")
        
        driver.close()
        
    except Exception as e:
        print(f"   Error: {e}")
    
    print("\n" + "=" * 60)
    print("✅ Ontology design demonstration complete!")
    print("=" * 60)
